"""
PatchGAN判别器模块
从SepMark项目移植的PatchGAN判别器实现
"""

import torch
import torch.nn as nn
from torch.nn import init
import functools

class PatchDiscriminator(nn.Module):
    """PatchGAN判别器"""
    
    def __init__(self, input_nc=3, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
        """
        初始化PatchGAN判别器
        
        Args:
            input_nc: 输入通道数
            ndf: 第一层的特征数
            n_layers: 判别器层数
            norm_layer: 归一化层类型
        """
        super(PatchDiscriminator, self).__init__()
        
        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d

        kw = 4
        padw = 1
        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), 
                   nn.LeakyReLU(0.2, True)]
        
        nf_mult = 1
        nf_mult_prev = 1
        for n in range(1, n_layers):
            nf_mult_prev = nf_mult
            nf_mult = min(2 ** n, 8)
            sequence += [
                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
                norm_layer(ndf * nf_mult),
                nn.LeakyReLU(0.2, True)
            ]

        nf_mult_prev = nf_mult
        nf_mult = min(2 ** n_layers, 8)
        sequence += [
            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
            norm_layer(ndf * nf_mult),
            nn.LeakyReLU(0.2, True)
        ]

        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
        self.model = nn.Sequential(*sequence)

    def forward(self, input):
        """前向传播"""
        return self.model(input)

class MultiscaleDiscriminator(nn.Module):
    """多尺度判别器"""
    
    def __init__(self, input_nc=3, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, num_D=3):
        """
        初始化多尺度判别器
        
        Args:
            input_nc: 输入通道数
            ndf: 第一层特征数
            n_layers: 每个判别器的层数
            norm_layer: 归一化层类型
            num_D: 判别器数量
        """
        super(MultiscaleDiscriminator, self).__init__()
        self.num_D = num_D
        self.n_layers = n_layers
        
        for i in range(num_D):
            netD = PatchDiscriminator(input_nc, ndf, n_layers, norm_layer)
            setattr(self, 'discriminator_%d' % i, netD)

        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)

    def forward(self, input):
        """前向传播"""
        result = []
        get_intermediate_features = not self.training
        
        for name, D in self.named_children():
            if 'discriminator' in name:
                out = D(input)
                if not get_intermediate_features:
                    result.append(out)
                else:
                    result.append([out])
            input = self.downsample(input)
        return result

def init_weights(net, init_type='normal', init_gain=0.02):
    """初始化网络权重"""
    def init_func(m):
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)

def define_discriminator(input_nc=3, ndf=64, netD='basic', n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
    """定义判别器"""
    net = None
    norm_layer = get_norm_layer(norm_type=norm)

    if netD == 'basic':  # default PatchGAN classifier
        net = PatchDiscriminator(input_nc, ndf, n_layers=3)
    elif netD == 'n_layers':  # more options
        net = PatchDiscriminator(input_nc, ndf, n_layers_D)
    elif netD == 'multiscale':  # multiscale discriminator
        net = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, num_D=3)
    else:
        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
    
    return init_net(net, init_type, init_gain, gpu_ids)

def get_norm_layer(norm_type='instance'):
    """获取归一化层"""
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        def norm_layer(x): 
            return nn.Identity()
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer

def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """初始化网络"""
    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.to(gpu_ids[0])
        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
    init_weights(net, init_type, init_gain=init_gain)
    return net

# RaLSGAN损失函数
class RaLSGANLoss(nn.Module):
    """相对平均最小二乘GAN损失"""
    
    def __init__(self):
        super(RaLSGANLoss, self).__init__()
        self.mse_loss = nn.MSELoss()
    
    def forward(self, pred_real, pred_fake, is_discriminator=True):
        """
        计算RaLSGAN损失
        
        Args:
            pred_real: 真实样本的判别器输出
            pred_fake: 假样本的判别器输出  
            is_discriminator: 是否为判别器损失
        """
        if is_discriminator:
            # 判别器损失
            real_logits = pred_real - torch.mean(pred_fake)
            fake_logits = pred_fake - torch.mean(pred_real)
            
            loss_real = self.mse_loss(real_logits, torch.ones_like(real_logits))
            loss_fake = self.mse_loss(fake_logits, -torch.ones_like(fake_logits))
            
            return (loss_real + loss_fake) / 2
        else:
            # 生成器损失
            real_logits = pred_real - torch.mean(pred_fake)
            fake_logits = pred_fake - torch.mean(pred_real)
            
            loss_real = self.mse_loss(real_logits, -torch.ones_like(real_logits))
            loss_fake = self.mse_loss(fake_logits, torch.ones_like(fake_logits))
            
            return (loss_real + loss_fake) / 2

def test_patch_discriminator():
    """测试PatchGAN判别器"""
    print("Testing PatchGAN Discriminator...")
    
    # 创建判别器
    discriminator = PatchDiscriminator(input_nc=3, ndf=64, n_layers=3)
    
    # 创建测试输入
    test_input = torch.randn(2, 3, 256, 256)
    
    # 前向传播
    output = discriminator(test_input)
    
    print(f"Input shape: {test_input.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Parameters: {sum(p.numel() for p in discriminator.parameters())}")
    
    # 测试RaLSGAN损失
    print("\nTesting RaLSGAN Loss...")
    criterion = RaLSGANLoss()
    
    pred_real = torch.randn(2, 1, 8, 8)
    pred_fake = torch.randn(2, 1, 8, 8)
    
    d_loss = criterion(pred_real, pred_fake, is_discriminator=True)
    g_loss = criterion(pred_real, pred_fake, is_discriminator=False)
    
    print(f"Discriminator loss: {d_loss.item():.4f}")
    print(f"Generator loss: {g_loss.item():.4f}")
    
    print("Test completed successfully!")

if __name__ == "__main__":
    test_patch_discriminator() 